#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import serial

# 导入达妙电机控制模块
try:
    from DM_CAN import Motor, MotorControl, DM_Motor_Type, Control_Type
    DM_CAN_AVAILABLE = True
    print("✅ DM_CAN电机控制模块导入成功")
except Exception as e:
    DM_CAN_AVAILABLE = False
    print(f"⚠️ DM_CAN电机控制模块导入失败: {e}")


class Robotarm:
    """机械臂控制类：读取限位配置并通过DM_CAN的MIT模式设定位置"""

    def __init__(self, config_path: str = "/home/tz2/T9-Test/arm_config.json"):
        if not DM_CAN_AVAILABLE:
            raise RuntimeError("DM_CAN模块不可用，无法初始化Robotarm")

        # 读取配置
        self.config_path = config_path
        self.config = self._load_config(self.config_path)
        self.arm_side = self.config.get("arm_side", "left")

        # 串口初始化
        serial_cfg = self.config.get("serial", {})
        port = serial_cfg.get("port", "/dev/ttyACM0")
        baudrate = serial_cfg.get("baudrate", 921600)
        timeout = serial_cfg.get("timeout", 0.5)

        self.serial = serial.Serial(port=port, baudrate=baudrate, timeout=timeout)
        self.mc = MotorControl(self.serial)

        # 电机类型字符串到枚举映射
        self._type_map = {
            "DM4310": DM_Motor_Type.DM4310,
            "DM4310_48V": DM_Motor_Type.DM4310_48V,
            "DM4340": DM_Motor_Type.DM4340,
            "DM4340_48V": DM_Motor_Type.DM4340_48V,
            "DM6006": DM_Motor_Type.DM6006,
            "DM8006": DM_Motor_Type.DM8006,
            "DM8009": DM_Motor_Type.DM8009,
            "DM10010L": DM_Motor_Type.DM10010L,
            "DM10010": DM_Motor_Type.DM10010,
            "DMH3510": DM_Motor_Type.DMH3510,
            "DMH6215": DM_Motor_Type.DMH6215,
            "DMG6220": DM_Motor_Type.DMG6220,
        }

        # 注册电机并建立ID映射
        self.motors_by_id = {}
        for m in self.config.get("motors", []):
            m_type_str = m.get("type", "DM4310")
            m_type = self._type_map.get(m_type_str, DM_Motor_Type.DM4310)
            slave_id = int(m.get("slave_id"))
            master_id = int(m.get("master_id", 0))
            motor = Motor(m_type, slave_id, master_id)
            self.mc.addMotor(motor)
            self.motors_by_id[slave_id] = motor

        # 读取位置限位
        self.position_limits_by_id = {}
        self._load_position_limits()

        # 切换到MIT模式（若固件支持）
        for motor in self.motors_by_id.values():
            try:
                self.mc.switchControlMode(motor, Control_Type.MIT)
            except Exception:
                # 保持静默，旧固件可能不支持
                pass

    def _load_config(self, path: str) -> dict:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)

    def _parse_id(self, id_str: str) -> int:
        try:
            if isinstance(id_str, int):
                return id_str
            id_str = str(id_str)
            if id_str.startswith("0x") or id_str.startswith("0X"):
                return int(id_str, 16)
            return int(id_str)
        except Exception:
            raise ValueError(f"电机ID解析失败: {id_str}")

    def get_arm_torque(self):
        """获取机械臂所有电机的力矩
        
        Returns:
            dict: 电机ID到力矩值的映射，格式为 {slave_id: torque_value}
        """
        torque_data = {}
        
        # 先接收一次数据以更新电机状态
        try:
            self.mc.recv()
        except Exception as e:
            print(f"⚠️ 接收电机数据时出错: {e}")
        
        # 遍历所有注册的电机，获取力矩数据
        for slave_id, motor in self.motors_by_id.items():
            try:
                torque = motor.getTorque()
                torque_data[slave_id] = torque
            except Exception as e:
                print(f"⚠️ 获取电机{slave_id}力矩时出错: {e}")
                torque_data[slave_id] = 0.0
        
        return torque_data

    def get_arm_temp(self):
        """获取机械臂所有电机的温度
        
        Returns:
            dict: 电机ID到温度值的映射，格式为 {slave_id: {'mos_temp': float, 'rotor_temp': float}}
        """
        temp_data = {}
        
        # 先接收一次数据以更新电机状态
        try:
            self.mc.recv()
        except Exception as e:
            print(f"⚠️ 接收电机数据时出错: {e}")
        
        # 遍历所有注册的电机，获取温度数据
        for slave_id, motor in self.motors_by_id.items():
            try:
                mos_temp, rotor_temp = motor.getTemperature()
                temp_data[slave_id] = {
                    'mos_temp': mos_temp,
                    'rotor_temp': rotor_temp
                }
            except Exception as e:
                print(f"⚠️ 获取电机{slave_id}温度时出错: {e}")
                temp_data[slave_id] = {
                    'mos_temp': 0.0,
                    'rotor_temp': 0.0
                }
        
        return temp_data

    def _load_position_limits(self):
        side_limits = self.config.get("motor_limits", {}).get(self.arm_side, {})
        for id_str, lim in side_limits.items():
            slave_id = self._parse_id(id_str)
            min_val = lim.get("min")
            max_val = lim.get("max")
            self.position_limits_by_id[slave_id] = (min_val, max_val)

    def _clamp_position(self, slave_id: int, pos: float) -> float:
        lim = self.position_limits_by_id.get(slave_id)
        if not lim:
            return pos
        min_val, max_val = lim
        if min_val is not None and pos < min_val:
            return min_val
        if max_val is not None and pos > max_val:
            return max_val
        return pos

    def set_pos(self, slave_id: int, position: float, kp: float = 50.0, kd: float = 1.0,
                dq: float = 0.0, tau: float = 0.0) -> float:
        """设置指定电机位置（MIT模式），自动应用限位

        Args:
            slave_id: 电机从机ID
            position: 目标位置（rad）
            kp: 位置环KP（0~500，DM固件范围）
            kd: 位置环KD（0~5，DM固件范围）
            dq: 期望速度（rad/s）
            tau: 期望力矩

        Returns:
            实际发送的目标位置（经过限位夹紧后）
        """
        if slave_id not in self.motors_by_id:
            raise KeyError(f"未找到电机，SlaveID={slave_id}")

        motor = self.motors_by_id[slave_id]

        # 限位夹紧
        q_target = self._clamp_position(slave_id, position)

        # 发送MIT控制
        self.mc.controlMIT(motor, kp, kd, q_target, dq, tau)
        return q_target

    def Armcontroller(self):
        """机械臂控制器示例"""  
        def set_single_joint_pos(slave_id,angle,kp,kd):
            angle = position
            set_pos(slave_id,position,kp,kd)
        def move_arm_x():
            pass
        def move_arm_y():
            pass
        def move_arm_z():
            pass
        def move_arm_pitch():
            pass
        def move_arm_roll():
            pass
        def move_arm_yaw():
            pass
        def stop():
            """停止机械臂运动"""
            set_pos(self.arm_motor_id, 0.0, kp=0.0, kd=0.0, dq=0.0, tau=0.0)
            pass
        pass